#ifndef	gema_included
#define	gema_included

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
struct	gema {
    static	tensor<half,S*H*C>	gi;
    static	tensor<half,2*S*H*C>	gx;
    tensor<half,R*C>	in;
    tensor<half,2*S*H*C>	x0,xo;
    tensor<half,S*H*C>	tmp;
    rmsnorm<2*S*H,C,2*H>	n;
    linear<R,2*S*H,C>	x;
    ema<2*S*H,C>	m;
    gating<S*H,C>	g;
    linear<S*H,R,C>	o;

    void	save(FILE	*F) {
        x.save(F);
        m.save(F);
        o.save(F);
    }

    void	load(FILE	*F) {
        x.load(F);
        m.load(F);
        o.load(F);
    }

    uint64_t	size(void) {
        return	x.size()+m.size()+o.size();
    }

    void	fw(tensor<half,R*C>	&inp) {
        cudaMemcpy(in.p,inp.p,R*C*sizeof(half),cudaMemcpyDeviceToDevice);
        x.fwt(in,x0,0);
        m.fw(x0,xo);
        n.fw(xo);
        g.fw(xo,tmp);
        o.fw(tmp,inp,resnet_scaling);
    }

    void	bk(tensor<half,R*C>	&gin) {
        o.bk(tmp,gin,gi,0);
        g.bk(xo,gi,gx);
        n.bk(xo,gx);
        m.bk(x0,gx,xo);
        x.bk(in,xo,gin,resnet_scaling);
    }
};

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
tensor<half,S*H*C>	gema<R,C,H,S>::gi;
template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
tensor<half,2*S*H*C>	gema<R,C,H,S>::gx;
#endif

